import numpy as np
import cv2
from matplotlib import pyplot as plt
import torch
import torch.nn.functional as F
from utils.metrix import findPeakPoint
from LMdata.LMdataset import idx2kp

colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
          [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
          [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
          ]



def ImgTensorToNumpy(img):
    '''
    :param img: torch FloatTensor of shape (3,H,W)  RGB  or (H,W)  Heatmap
    :return:  numpy of shape (H,W,3)  RGB  or (H,W)
    '''
    if len(img.size())==3:
        return np.uint8(img.numpy().clip(0.,255.)).transpose((1,2,0))
    return img.numpy().clip(0., 255.)

def ImgNumpyToTensor(img):
    '''
    :param img: numpy of shape (H,W,3) RGB or (H,W)
    :return: torch FloatTensor of shape (3,H,W) RGB or (H,W)
    '''
    if len(img.shape)==3:
        return torch.from_numpy(img.transpose((2,0,1))).float()
    return torch.from_numpy(img).float()


def AddHeatmap(img, heatmap):
    '''
    :param img:  np.array  (H,W,3) RGB
    :param heatmap: np.array  (H,W)  OK to be not normalized
    :return:
    '''
    # normalize heatmap
    heatmap = np.maximum(heatmap, 0)
    heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap))  # Normalize between 0-1
    heatmap = np.uint8(heatmap * 255)  # Scale between 0-255 to visualize

    heatmap = np.repeat(heatmap[:, :, np.newaxis], 3, axis=2)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_HSV)

    img_with_heatmap = np.float32(heatmap) + np.float32(img)

    img_with_heatmap = img_with_heatmap / np.max(img_with_heatmap)
    img_with_heatmap = np.uint8(255 * img_with_heatmap)
    return img_with_heatmap

def DrawKeyPoints(img, kp):
    '''
    :param img:  np.array  (H,W,3) RGB
    :param kp:  np.array  (N,3) x,y,vis  %
    :return:
    '''
    h,w = img.shape[0:2]

    for i,(x,y,vis) in enumerate(kp):
        if vis>-1 and x>=0 and y>=0:
            x = int(x*w)
            y = int(y*h)
            if vis==1:
                cv2.circle(img, (x, y), 6, colors[i%len(colors)], thickness=-1)

            else:
                cv2.circle(img, (x, y), 6, colors[i%len(colors)], thickness= 2)

    return img



def getPlotImg(smp_img, pred_hm, true_hm, flm):
    H,W = smp_img.size()[-2:]

    imgs_to_plot = np.zeros((4, H, W, 3))

    img_mean = torch.from_numpy(np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)).float().cuda()
    img_std = torch.from_numpy(np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)).float().cuda()

    ori_img = (smp_img.data*img_std+img_mean)*255.
    ori_img = ImgTensorToNumpy(ori_img.cpu())

    pred_hm = pred_hm.unsqueeze(0)
    true_hm = true_hm.unsqueeze(0)


    if pred_hm.size()[-2:] != smp_img.size()[-2:]:
        pred_hm = F.upsample(pred_hm, size=smp_img.size()[-2:], mode='bilinear')  # (1,24,336,336)
        true_hm = F.upsample(true_hm, size=smp_img.size()[-2:], mode='bilinear')  # (1,24,336,336)

    pred_kp = np.zeros((24, 3))
    pred_kp[:, 0:2] = findPeakPoint(pred_hm.data.cpu().numpy())[0]
    pred_kp[:, 2] = flm[:, 2]

    # draw keypoints
    imgs_to_plot[2] = DrawKeyPoints(ori_img.copy(), flm)
    imgs_to_plot[3] = DrawKeyPoints(ori_img.copy(), pred_kp)

    pred_hm = pred_hm[0].sum(dim=0)
    true_hm = true_hm[0].sum(dim=0)

    # add heatmap to origin image
    imgs_to_plot[0] = AddHeatmap(ori_img, ImgTensorToNumpy(true_hm.data.cpu()))
    imgs_to_plot[1] = AddHeatmap(ori_img, ImgTensorToNumpy(pred_hm.data.cpu()))



    return imgs_to_plot



if __name__ == '__main__':
    from preprocessing import get_label_from_df
    from LMdata.LMdataset import kpInCates, idx2cate,cate2idx, idx2kp, kp2idx
    import pandas as pd
    import os
    # df = pd.read_csv('/home/gserver/zhangchi/Landmarks/val1/GCN336_bs4_5.85_336pred.csv')

    df = pd.read_csv('/media/hszc/data1/FaishionAI/round2_crop/val/Annotations/GCN512_bs16_3.98_attr_val.csv')


    cates = df.image_category.apply(lambda x:cate2idx[x]).tolist()
    paths =  df.image_id.tolist()
    print paths[0]
    # paths = df.image_id.apply(lambda x:os.path.join('/media/gserver/data/landmark/rawdata/test_b',x)).tolist()

    # for cate_idx in idx2cate:
        # cate = idx2cate[cate_idx]
        # for kp in kp2idx:
        #     kp_idx = kp2idx[kp]
        #     if kp_idx not in kpInCates[cate_idx]:
        #
        #         df.loc[df['image_category']==cate, kp] = '-1_-1_-1'



    pred_xy = get_label_from_df(df).astype(float)
    for i,kps in enumerate(pred_xy):
        img = cv2.imread(paths[i])
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        kps[:,0] = 1. * kps[:,0] / img.shape[1]
        kps[:,1] = 1. * kps[:,1] / img.shape[0]

        img2 = DrawKeyPoints(img,kps)
        plt.imshow(img2)
        plt.show()
    #
